import json
import logging
import os
import sys
from typing import Optional
from staintools import StainNormalizer
from PIL import Image
from openslide import OpenSlide
from shapely import MultiPolygon
from shapely.geometry import Polygon, Point, shape
import numpy as np
import cv2
from util.file_util import get_svs_files
from util.log_util import get_logger
from tqdm import tqdm

# 标准图片
targetPath = "../static/target.jpg"


# 示例：判断 patch 是否落在病灶区域
def is_patch_inside_roi(patch_x: float, patch_y: float, patch_size: int, polygon: Optional[MultiPolygon]):
    center = Point(patch_x + patch_size // 2, patch_y + patch_size // 2)
    return polygon.contains(center)


# patch_size 是每个切图块的尺寸, stride 是滑动的距离。
def getAllPatchs(patch_size: int, stride: int, pathSVS: str, polygon: Optional[MultiPolygon], r: bool) -> list[Image]:
    # 加载SVS
    slide = OpenSlide(pathSVS)
    width, height = slide.dimensions
    patches = []
    for y in range(0, height, stride):
        for x in range(0, width, stride):
            if is_patch_inside_roi(x, y, patch_size, polygon) == r:
                patch = slide.read_region((x, y), level=0, size=(patch_size, patch_size))
                # 过滤没有组织的图像
                if is_background(patch):
                    continue
                patches.append(patch)
                # 提取反向数据
                if len(patches) == 2000 and r == False:
                    # 如果是父类图片，只提取2000张即可！！！
                    return patches
                if len(patches) % 1000 == 0:
                    logging.info("%s 完成 %d patchs提取。", pathSVS, len(patches))
    return patches


# 根据geojson提取多个区域的病灶
def getAllPointJsonData(path: str) -> Optional[Polygon]:
    # 加载 .geojson 文件
    with open(path, "r") as f:
        geojson_data = json.load(f)
    # 提取 polygon（支持多个区域）
    polygons = []
    for feature in geojson_data["features"]:
        geometry = feature["geometry"]
        if geometry["type"] == "Polygon":
            poly = shape(geometry)  # 转换为 shapely Polygon 对象
            polygons.append(poly)
    return polygons


def saveAllPatch(patchs: list[Image], path: str):
    output_dir = path
    os.makedirs(output_dir, exist_ok=True)  # 如果文件夹不存在就创建
    for i, patch in enumerate(patchs):
        save_path = os.path.join(output_dir, f"patch_{i:04d}.png")
        patch.save(save_path)
        if i % 500 == 0:
            logging.info("已经成功保存%d张图片。", i)
    logging.info("保存成功，位置: %s, 数量: %d", output_dir, len(patchs))


# 是否为背景  true代表背景图  过滤没有组织的patch
def is_background(patch, threshold=220, coverage=0.95):
    """
    判断 patch 是否为背景。
    threshold: 像素值阈值，越高越偏白；
    coverage: 大于阈值的像素比例超过该值则认为是背景。
    """
    gray = patch.convert("L")
    np_gray = np.array(gray)
    white_ratio = np.mean(np_gray > threshold)
    return white_ratio > coverage


# colorNormalizer 颜色归一化
def colorNormalizer(patch: list[Image], targetPicture: str) -> list[Image]:
    # 使用一张参考 HE 图作为 target（最好是你预先选好的一张颜色标准图）
    target = cv2.imread(targetPicture)
    target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)
    normalizer = StainNormalizer(method='macenko')  # 可选：reinhard, macenko, vahadane
    normalizer.fit(target)

    normalized_patches = []
    for i, image in enumerate(tqdm(patch, desc="颜色标准化中")):
        # 过滤没有组织的图像
        if is_background(image):
            continue
        try:
            # 去除透明通道，强制转为 RGB
            img_rgb = np.array(image.convert("RGB"))
            # 异常数组检测(判断img_rgb是否是np.ndarray数组、img_rgb是否是三维度高宽通道、判断通道是否是3即rgb通道)
            if not isinstance(img_rgb, np.ndarray) or img_rgb.ndim != 3 or img_rgb.shape[2] != 3:
                logging.warning("跳过无效或非 RGB 图像：patch_%d，shape: %s", i, img_rgb.shape if isinstance(img_rgb, np.ndarray) else "Invalid")
                continue
            # 补充组织面积筛查：暗色像素过少（说明没有组织）
            # img_rgb < 220：比较图像中每个像素的 RGB 值是否小于 220（即是否是比较暗的颜色），返回一个布尔数组。 RGB 值越小，颜色越暗
            # np.mean(img_rgb < 220)：把布尔数组转为 0/1，并计算平均值，代表“暗色像素占总像素比例”。
            # 如果暗色像素比例低于 5%，说明这个 patch 几乎是亮的（很可能是背景、玻片边缘等），没有组织。跳过处理。
            tissue_ratio = np.mean(img_rgb < 220)
            if tissue_ratio < 0.05:
                logging.warning("跳过组织面积过小的 patch_%d", i)
                continue
            img_norm = normalizer.transform(img_rgb)
            img_norm = np.clip(img_norm, 0, 255).astype(np.uint8)
            normalized_patches.append(Image.fromarray(img_norm))
        except Exception as e:
            logging.error("标准化失败：patch_%d，错误: %s", i, str(e))
    return normalized_patches


def getValidData(path: str, fileName: str):
    # 初始化日志格式
    get_logger(logging.INFO)
    # 第一步: 得到病灶区域坐标
    jsonPath = os.path.splitext(name)[0] + ".geojson"
    data = getAllPointJsonData(os.path.join(path, jsonPath))
    polygon = MultiPolygon(data)
    for i, p in enumerate(data):
        logging.info("%s 病灶区域%d数据点数量: %d", fileName, i + 1, len(p.exterior.coords))
    # 第二步: 得到病灶区域的patch
    logging.info("提取patchs中......")
    svsPath = os.path.splitext(name)[0] + ".svs"
    patchs = getAllPatchs(512, 512, os.path.join(path, svsPath), polygon,
                          True)
    logging.info("切割病灶(未标准化)得到的patch数量: %d", len(patchs))
    # 第三步：空白区域检测 + 颜色标准化
    logging.info("%s 的patch开始归一化...", fileName)
    normalizer = colorNormalizer(patchs, targetPath)
    logging.info("%s 的patch归一化完成", fileName)
    # 第四步：保存图片
    path1 = os.path.join(path, os.path.splitext(name)[0])
    saveAllPatch(normalizer, os.path.join(path1, "valid"))
    logging.info("%s 保存图片完成！  paths数量: %d", fileName, len(normalizer))


# path 有svs和geojson文件的路径，fileName svs的文件名称
def getInValidData(path: str, fileName: str):
    # 初始化日志格式
    get_logger(logging.INFO)
    # 第一步: 得到病灶区域坐标
    jsonPath = os.path.splitext(name)[0] + ".geojson"
    data = getAllPointJsonData(os.path.join(path, jsonPath))
    polygon = MultiPolygon(data)
    for i, p in enumerate(data):
        logging.info("%s 病灶区域%d数据点数量: %d", fileName, i + 1, len(p.exterior.coords))
    # 第二步: 得到病灶区域的patch
    logging.info("提取patchs中......")
    svsPath = os.path.splitext(name)[0] + ".svs"
    patchs = getAllPatchs(512, 512, os.path.join(path, svsPath), polygon,
                          False)
    logging.info("切割病灶(未标准化)得到的patch数量: %d", len(patchs))
    # 第三步：空白区域检测 + 颜色标准化
    logging.info("%s 的patch开始归一化...", fileName)
    normalizer = colorNormalizer(patchs, targetPath)
    logging.info("%s 的patch归一化完成", fileName)
    # 第四步保存图像
    path1 = os.path.join(path, os.path.splitext(name)[0])
    saveAllPatch(normalizer, os.path.join(path1, "invalid"))


if __name__ == '__main__':
    # 1. svs存放的地址
    svsPath = "D:/zm_scientific_research/supervised_cutting"
    # 2. 获取所有的svs文件名称
    files = get_svs_files(svsPath)
    for name, path in files.items():
        print(f"{name} -> {path}")
        logging.info("%s --  提取正向path", name)
        getValidData(svsPath, name)
        logging.info("%s --  提取反向path", name)
        getInValidData(svsPath, name)
